package dd.lo.pdfimg;

import dd.lo.HelloCV;
import dd.lo.util.ImageUtils;
import javafx.application.Application;
import javafx.application.Platform;
import javafx.concurrent.Task;
import javafx.scene.Scene;
import javafx.scene.control.Button;
import javafx.scene.image.Image;
import javafx.scene.image.ImageView;
import javafx.scene.layout.BorderPane;
import javafx.scene.layout.FlowPane;
import javafx.stage.FileChooser;
import javafx.stage.Stage;
import org.opencv.core.*;
import org.opencv.imgcodecs.Imgcodecs;
import org.opencv.imgproc.Imgproc;

import java.io.ByteArrayInputStream;
import java.io.File;
import java.util.ArrayList;
import java.util.List;

public class ImageCropper extends Application {

    public static void main(String[] args) {
        System.load(HelloCV.class.getResource("/libopencv_java481.dylib").getPath());
        ImageCropper.start(args);
    }
    public static void start(String[] args) {
        launch(args);
    }

    private static final int DEFAULT_WIDTH = 1080;
    private static final int DEFAULT_HEIGHT = 768;

    private ImageView imageView;

    private int mode = 0;

    public static Mat[] cropImage(Mat src) {
        Mat grayImg = new Mat();
        Imgproc.cvtColor(src, grayImg, Imgproc.COLOR_BGR2GRAY);
        //Preprocess
        Core.normalize(grayImg, grayImg, 0, 255, Core.NORM_MINMAX);
        //Use threshold
        Mat bw = new Mat();
        Imgproc.threshold(grayImg, bw, 240, 255, Imgproc.THRESH_BINARY);
        //Find small images
        List<MatOfPoint> contours = new ArrayList<>();
        Mat hierarchy = new Mat();
        Imgproc.findContours(bw, contours, hierarchy, Imgproc.RETR_TREE, Imgproc.CHAIN_APPROX_SIMPLE);
        List<Rect> rois = new ArrayList<>(contours.size());
        int minContourWidth = (int) Math.round(0.05 * src.width());
        int maxContourWidth = (int) Math.round(0.8 * src.width());
        int minContourHeight = (int) Math.round(0.05 * src.height());
        int maxContourHeight = (int) Math.round(0.8 * src.height());
        for (int i = 0; i < contours.size(); ++i) {
            MatOfPoint contour = contours.get(i);
            Rect rect = Imgproc.boundingRect(contour);
//            if (rect.width > maxContourWidth || rect.height > maxContourHeight || rect.width < minContourWidth || rect.height < minContourHeight) {
            if (rect.width > maxContourWidth || rect.height > maxContourHeight || rect.width < minContourWidth) {
                continue;
            }
            //0~3分别表示后一个轮廓、前一个轮廓、父轮廓、内嵌轮廓的索引值
            double[] h = hierarchy.get(0, i);
            if (h[3] == 0.0) {
                //如果这个轮廓是最外围的，则内嵌轮廓的索引值为0
                //Store contour boundingRect
                rois.add(rect);
            }
        }
        Mat[] result = new Mat[rois.size()];
        for (int i = 0; i < rois.size(); ++i) {
            result[i] = new Mat(src, rois.get(i));
        }
        return result;
    }

    private void cropImg(Mat src) {
        Task<Void> task = new Task<Void>() {
            @Override
            protected Void call() {
                Mat grayImg = new Mat();
                Imgproc.cvtColor(src, grayImg, Imgproc.COLOR_BGR2GRAY);
                //Preprocess
//                Imgproc.GaussianBlur(grayImg, grayImg, new Size(3, 3), 0);
//                Imgproc.equalizeHist(grayImg, grayImg);
                Core.normalize(grayImg, grayImg, 0, 255, Core.NORM_MINMAX);
                //Use threshold
                Mat bw = new Mat();
                Imgproc.threshold(grayImg, bw, 220, 255, Imgproc.THRESH_BINARY);
                //锐化图像
                Mat kernel = new Mat(3, 3, CvType.CV_32SC1);
//                kernel.put(0, 0, new int[]{0, -1,  0, -1,  5, -1, 0, -1,  0});
                kernel.put(0, 0, new int[]{-1, -1,  -1, -1,  9, -1, -1, -1,  -1});
                Imgproc.filter2D(bw, bw, bw.depth(), kernel);
                //给图像加上外围的白边，让涂黑的轮廓不会粘到图片边缘导致不能被识别
                Core.copyMakeBorder(bw, bw, 10, 10, 10, 10, Core.BORDER_CONSTANT, new Scalar(255, 255, 255));
                if (mode == 1) {
                    //这里只是查看前面处理的结果
                    ImageUtils.fitImgSize(bw, DEFAULT_WIDTH, DEFAULT_HEIGHT);
                    MatOfByte buffer = new MatOfByte();
                    Imgcodecs.imencode(".jpeg", bw, buffer);
                    Platform.runLater(() -> imageView.setImage(new Image(new ByteArrayInputStream(buffer.toArray()))));
                    return null;
                }
                //Find small images
                //参考资料：https://answers.opencv.org/question/86328/search-small-images-on-one-images/
                List<MatOfPoint> contours = new ArrayList<>();
                Mat hierarchy = new Mat();
                Imgproc.findContours(bw, contours, hierarchy, Imgproc.RETR_TREE, Imgproc.CHAIN_APPROX_SIMPLE);
                List<Rect> rois = new ArrayList<>(contours.size());
                int minContourWidth = (int) Math.round(0.15 * src.width());
                int maxContourWidth = (int) Math.round(0.8 * src.width());
                int minContourHeight = (int) Math.round(0.15 * src.height());
                int maxContourHeight = (int) Math.round(0.8 * src.height());
                for (int i = 0; i < contours.size(); ++i) {
                    MatOfPoint contour = contours.get(i);
                    Rect rect = Imgproc.boundingRect(contour);
                    //因为加了白边，把白边的位移补偿回去
                    rect.x -= 10;
                    rect.y -= 10;
//                    if (rect.width > maxContourWidth || rect.height > maxContourHeight) {
//                        continue;
//                    }
                    if (rect.width < minContourWidth || rect.height < minContourHeight){
                        if (minContourWidth / 2.0 < rect.width && minContourHeight / 2.0 < rect.height) {
                            //Draw not so small contours in RED
                            Imgproc.drawContours(src, contours, i, new Scalar(0, 0, 255), 1);
                        }
                        continue;
                    }
                    //0~3分别表示后一个轮廓、前一个轮廓、父轮廓、内嵌轮廓的索引值
                    //参考资料：https://blog.csdn.net/yuanren201/article/details/104418710
                    double[] h = hierarchy.get(0, i);
                    if (h[3] == 0.0) {
                        //如果这个轮廓是最外围的，则内嵌轮廓的索引值为0
                        //Draw selected contours in GREEN
                        Imgproc.drawContours(src, contours, i, new Scalar(0, 255, 0), 1, Imgproc.LINE_4, hierarchy, 0);
                        //Store contour boundingRect and draw it in BLUE
                        Imgproc.rectangle(src, rect, new Scalar(255, 0, 0), 4);
                        rois.add(rect);
                    }
                }
                System.out.printf("Find %d ROIs\n", rois.size());
                ImageUtils.fitImgSize(src, DEFAULT_WIDTH, DEFAULT_HEIGHT);
                MatOfByte buffer = new MatOfByte();
                Imgcodecs.imencode(".jpeg", src, buffer);
                Platform.runLater(() -> imageView.setImage(new Image(new ByteArrayInputStream(buffer.toArray()))));
                return null;
            }
        };
        new Thread(task).start();
    }

    @Override
    public void start(Stage stage) {
        imageView = new ImageView();
        //File chooser
        Button imageChooser = new Button("选择图片");
        imageChooser.setOnMouseClicked(e -> {
            FileChooser fileChooser = new FileChooser();
            fileChooser.setTitle("选择图片");
            fileChooser.setInitialDirectory(new File(System.getProperty("user.home") + "/Downloads/train"));
            fileChooser.getExtensionFilters().addAll(
                    new FileChooser.ExtensionFilter("JPEG", "*.jpeg")
            );
            File imgFile = fileChooser.showOpenDialog(stage);
            if (null == imgFile) {
                System.out.println("用户取消选中文件");
                return;
            }
            //读取一张图片
            Mat src = Imgcodecs.imread(imgFile.getAbsolutePath(), Imgcodecs.IMREAD_UNCHANGED);
            cropImg(src);
        });
        Button modeBtn = new Button("切换模式");
        modeBtn.setOnMouseClicked(event -> {
            ++mode;
            mode %= 2;
            System.out.printf("mod=%d\n", mode);
        });
        FlowPane bottomControl = new FlowPane();
        bottomControl.getChildren().addAll(imageChooser, modeBtn);
        bottomControl.setPrefHeight(50);
        BorderPane root = new BorderPane();
        root.setCenter(imageView);
        root.setBottom(bottomControl);
        Scene scene = new Scene(root, DEFAULT_WIDTH, DEFAULT_HEIGHT + 50);
        stage.setScene(scene);
        stage.setTitle("Image cropping example");
        stage.setResizable(false);
        stage.show();
    }
}
